import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from nicegui import ui
import numpy as np
from datetime import datetime
import time
import json
import os

# 全局变量存储原始数据
raw_df = None

# 用户数据管理
USER_DATA_FILE = 'user_data.json'

# 确保用户数据文件存在
if not os.path.exists(USER_DATA_FILE):
    with open(USER_DATA_FILE, 'w', encoding='utf-8') as f:
        json.dump({'admin': 'admin'}, f, ensure_ascii=False, indent=4)

# 加载用户数据
def load_users():
    try:
        with open(USER_DATA_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载用户数据失败: {e}")
        return {'admin': 'admin'}  # 默认管理员账户

# 保存用户数据
def save_users(users):
    try:
        with open(USER_DATA_FILE, 'w', encoding='utf-8') as f:
            json.dump(users, f, ensure_ascii=False, indent=4)
        return True
    except Exception as e:
        print(f"保存用户数据失败: {e}")
        return False

# 初始化用户数据
users = load_users()

# ------------------------------
# 登录页面
# ------------------------------
@ui.page('/')
def login_page():
    def handle_login():
        # 检查用户名和密码
        current_users = load_users()
        if username.value in current_users and current_users[username.value] == password.value:
            ui.navigate.to('/main')  # 登录成功跳转到主界面
        else:
            ui.notify("用户名或密码错误", type="negative")

    # 创建登录界面容器
    with ui.column().classes('absolute-center items-center w-full max-w-md p-4'):
        ui.label('股票数据分析平台').classes('text-3xl font-bold mb-8 text-blue-800')

        with ui.card().classes('w-full p-8 shadow-2xl'):
            ui.label('用户登录').classes('text-2xl font-bold mb-6 text-center')

            username = ui.input(label='用户名').classes('w-full mb-4')
            password = ui.input(label='密码', password=True).classes('w-full mb-6')

            ui.button('登录', on_click=handle_login).classes('w-full bg-blue-800 hover:bg-blue-700 text-white')

            # 注册链接
            with ui.row().classes('justify-center mt-4'):
                ui.label('还没有账号？')
                ui.button('立即注册', on_click=lambda: ui.navigate.to('/register')).classes('ml-2 bg-gray-600 hover:bg-gray-500 text-white')

# ------------------------------
# 注册页面
# ------------------------------
@ui.page('/register')
def register_page():
    def handle_register():
        # 验证表单
        if not username.value or not password.value:
            ui.notify("用户名和密码不能为空", type="negative")
            return
        
        if password.value != confirm_password.value:
            ui.notify("两次输入的密码不一致", type="negative")
            return
        
        # 检查密码强度
        if len(password.value) < 6:
            ui.notify("密码长度不能少于6位", type="negative")
            return
        
        current_users = load_users()
        if username.value in current_users:
            ui.notify("用户名已存在", type="negative")
            return
        
        # 添加新用户
        current_users[username.value] = password.value
        if save_users(current_users):
            ui.notify("注册成功，请登录", type="positive")
            ui.navigate.to('/')  # 注册成功后跳转到登录页面
        else:
            ui.notify("注册失败，请重试", type="negative")

    # 创建注册界面容器
    with ui.column().classes('absolute-center items-center w-full max-w-md p-4'):
        ui.label('股票数据分析平台').classes('text-3xl font-bold mb-8 text-blue-800')

        with ui.card().classes('w-full p-8 shadow-2xl'):
            ui.label('用户注册').classes('text-2xl font-bold mb-6 text-center')

            username = ui.input(label='用户名').classes('w-full mb-4')
            password = ui.input(label='密码', password=True).classes('w-full mb-4')
            confirm_password = ui.input(label='确认密码', password=True).classes('w-full mb-6')

            ui.button('注册', on_click=handle_register).classes('w-full bg-blue-800 hover:bg-blue-700 text-white')
            
            # 登录链接
            with ui.row().classes('justify-center mt-4'):
                ui.label('已有账号？')
                ui.button('返回登录', on_click=lambda: ui.navigate.to('/')).classes('ml-2 bg-gray-600 hover:bg-gray-500 text-white')

# ------------------------------
# 数据加载和处理函数
# ------------------------------
def load_and_preprocess():
    """加载并预处理数据"""
    global raw_df
    if raw_df is None:
        # 加载 CSV 文件
        raw_df = pd.read_csv('prices-split-adjusted.csv')
        # 处理缺失值
        raw_df = raw_df.dropna()
        # 转换日期格式
        raw_df['date'] = pd.to_datetime(raw_df['date'])
    return raw_df.copy()


def filter_data(df, start_date, end_date, selected_symbols):
    """根据筛选条件过滤数据"""
    # 日期筛选
    if start_date and end_date:
        try:
            start_dt = datetime.strptime(start_date, '%Y-%m-%d')
            end_dt = datetime.strptime(end_date, '%Y-%m-%d')
            df = df[(df['date'] >= start_dt) & (df['date'] <= end_dt)]
        except ValueError:
            ui.notify('日期格式错误，请使用 YYYY - MM - DD 格式', type='negative')

    # 股票代码筛选
    if selected_symbols and '所有代码' not in selected_symbols:
        df = df[df['symbol'].isin(selected_symbols)]

    # 如果股票代码什么都不选，返回空 DataFrame
    if not selected_symbols:
        return pd.DataFrame()

    return df

# ------------------------------
# 图表创建函数
# ------------------------------
def create_price_trend_chart(df):
    """价格趋势分析图"""
    if df.empty:
        return go.Figure()

    daily_price = df.groupby('date')['close'].mean().reset_index()

    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=daily_price['date'],
        y=daily_price['close'],
        name='每日收盘价',
        mode='lines',
        line=dict(color='#17becf', width=2),
        opacity=0.9
    ))

    fig.update_layout(
        title='每日收盘价趋势分析',
        xaxis_title='日期',
        yaxis_title='收盘价',
        hovermode='x unified',
        template='plotly_white',
        showlegend=True,
        height=500
    )
    return fig


def create_avg_close_price_bar_chart(df):
    """不同股票平均收盘价对比柱状图"""
    if df.empty:
        return go.Figure()

    symbol_avg_close = df.groupby('symbol')['close'].mean().reset_index()
    symbol_avg_close = symbol_avg_close.sort_values('close', ascending=False)

    fig = px.bar(
        symbol_avg_close,
        x='symbol',
        y='close',
        title='不同股票平均收盘价对比',
        labels={'close': '平均收盘价'}, 
        color='symbol',
        height=500
    )
    fig.update_layout(
        xaxis_title='股票代码',
        yaxis_title='平均收盘价',
        template='plotly_white',
        showlegend=False
    )
    return fig


def create_3d_scatter_chart(df):
    """3D 散点图"""
    if df.empty:
        return go.Figure()

    # 只取前 10 个股票代码避免图表过于拥挤
    top_symbols = df['symbol'].value_counts().head(10).index
    df_filtered = df[df['symbol'].isin(top_symbols)]

    # 对交易量取对数
    df_filtered['log_volume'] = np.log(df_filtered['volume'])

    fig = px.scatter_3d(
        df_filtered,
        x='high',
        y='close',
        z='log_volume',
        color='symbol',
        title='最高价、收盘价与对数交易量 3D 散点图',
        height=600
    )
    fig.update_layout(template='plotly_white')
    return fig


def create_volume_pie_chart(df):
    """不同股票成交量占比饼图"""
    if df.empty:
        return go.Figure()

    symbol_volume_sum = df.groupby('symbol')['volume'].sum().reset_index()
    total_volume = symbol_volume_sum['volume'].sum()
    symbol_volume_sum['volume_percentage'] = symbol_volume_sum['volume'].apply(
        lambda x: f'{(x / total_volume) * 100:.2f}%')

    fig = px.pie(
        symbol_volume_sum,
        values='volume',
        names='symbol',
        title='不同股票成交量占比',
        labels={'volume_percentage': '成交量占比'}, 
        hover_data=['volume_percentage'],
        height=500
    )
    fig.update_traces(textinfo='percent+label', textposition='inside')
    fig.update_layout(template='plotly_white')
    return fig

# ------------------------------
# 主界面
# ------------------------------
@ui.page('/main')
def main_page():
    # 初始化数据
    df = load_and_preprocess()
    all_symbols = sorted(df['symbol'].unique().tolist())

    # 当前选中的图表类型
    current_chart_type = '价格趋势'

    # 创建状态变量
    start_date = ui.input(
        label='开始日期',
        value=df['date'].min().strftime('%Y-%m-%d'),
        placeholder='YYYY-MM-DD'
    ).classes('w-full')

    end_date = ui.input(
        label='结束日期',
        value=df['date'].max().strftime('%Y-%m-%d'),
        placeholder='YYYY-MM-DD'
    ).classes('w-full')

    # 多选股票代码 - 初始状态为空
    selected_symbols = []
    # 添加特殊选项：所有代码和取消全选
    special_options = ['所有代码', '取消全选']

    # 处理选择变化
    def handle_selection_change(e):
        values = e.value

        # 处理"取消全选"选项
        if '取消全选' in values:
            # 清除所有选择
            symbols.set_value([])
            return

        # 处理"所有代码"选项
        if '所有代码' in values:
            # 如果选择了"所有代码"，则只显示"所有代码"而不显示具体股票代码
            symbols.set_value(['所有代码'])

    # 创建选择器
    symbols = ui.select(
        label='股票代码 (可多选)',
        options=special_options + all_symbols,
        value=selected_symbols,
        multiple=True,
    ).classes('w-full truncate-chips')

    # 添加CSS样式，确保选择框内选项过多时隐藏
    ui.add_css('''
        .truncate-chips .q-field__native {
            max-height: 40px;
            overflow-y: hidden;
            flex-wrap: nowrap;
        }
        .truncate-chips .q-chip {
            max-width: 100px;
            overflow: hidden;
            text-overflow: ellipsis;
        }
    ''')

    # 存储 KPI 的容器
    kpi_container = ui.row().classes('w-full max-w-7xl gap-4 justify-center')

    # 存储图表按钮的容器
    chart_buttons_container = ui.row().classes('w-full justify-center gap-4 mt-4')

    # 存储图表引用的容器
    chart_container = ui.column().classes('w-full mt-2')

    # 重置筛选条件的函数
    def reset_filters():
        start_date.value = df['date'].min().strftime('%Y-%m-%d')
        end_date.value = df['date'].max().strftime('%Y-%m-%d')
        symbols.set_value([])  # 重置为空选择
        ui.notify('筛选条件已重置!')
        update_visualizations()

    # 页面标题和筛选控件
    with ui.header().classes('bg-blue-800 text-white p-4 shadow-lg'):
        with ui.row().classes('items-center w-full justify-between'):
            # 标题
            ui.label('股票数据分析平台').classes('text-2xl font-bold')

            # 登出按钮
            ui.button('退出登录', on_click=lambda: ui.navigate.to('/')).classes('bg-red-500 hover:bg-red-600 text-white')

            # 筛选控件和操作按钮
            with ui.row().classes('items-center gap-4'):
                with ui.row().classes('items-center gap-4'):
                    with ui.column().classes('w-40'):
                        start_date.classes('bg-blue-700 text-white border-blue-600')

                    with ui.column().classes('w-40'):
                        end_date.classes('bg-blue-700 text-white border-blue-600')

                # 操作按钮
                with ui.row().classes('items-center gap-2'):
                    analyze_btn = ui.button('分析数据', on_click=lambda: update_visualizations())
                    analyze_btn.classes('bg-green-500 hover:bg-green-600 text-white py-2 px-4 rounded shadow')

                    reset_btn = ui.button('重置', on_click=reset_filters)
                    reset_btn.classes('bg-gray-600 hover:bg-gray-700 text-white py-2 px-4 rounded shadow')

    # 主内容区
    with ui.column().classes('w-full p-4 gap-6 bg-gray-100 min-h-screen items-center'):
        # KPI 指标卡容器
        kpi_container

        # 图表按钮容器（放在 KPI 指标卡下方）
        with chart_buttons_container:
            # 创建四个图表切换按钮
            price_btn = ui.button('价格趋势', on_click=lambda: switch_chart('价格趋势'))
            price_btn.classes('bg-blue-500 text-white hover:bg-blue-600')

            avg_btn = ui.button('平均收盘价', on_click=lambda: switch_chart('平均收盘价对比'))
            avg_btn.classes('bg-blue-500 text-white hover:bg-blue-600')

            scatter_btn = ui.button('3D 散点图', on_click=lambda: switch_chart('3D 散点图'))
            scatter_btn.classes('bg-blue-500 text-white hover:bg-blue-600')

            pie_btn = ui.button('成交量饼图', on_click=lambda: switch_chart('成交量饼图'))
            pie_btn.classes('bg-blue-500 text-white hover:bg-blue-600')

        # 图表显示容器
        chart_container

        # 初始状态显示提示信息
        with chart_container:
            ui.label('请选择股票代码并点击"分析数据"按钮').classes('text-xl text-gray-500 mt-8')

        # 原始数据表格容器紧贴在图表下方
        with ui.expansion('查看原始数据', icon='table_chart', value=False).classes(
                'w-full max-w-7xl bg-white rounded-lg shadow mt-4'):
            # 数据表格会在这里动态添加
            table_placeholder = ui.column().classes('w-full')

            # 添加居中显示的文字
            with ui.row().classes('w-full justify-center'):
                ui.label('数据量太大，这里显示其中 100 条数据')

    # 当前筛选后的数据
    current_filtered_df = None

    def switch_chart(chart_type):
        """切换显示不同的图表"""
        nonlocal current_chart_type
        current_chart_type = chart_type

        if current_filtered_df is None or current_filtered_df.empty:
            chart_container.clear()
            with chart_container:
                ui.label('请选择股票代码并点击"分析数据"按钮').classes('text-xl text-gray-500 mt-8')
            return

        # 更新按钮样式（突出显示当前选中的按钮）
        buttons = chart_buttons_container.default_slot.children
        for btn in buttons:
            btn.classes(remove='bg-blue-700', replace='bg-blue-500 text-white hover:bg-blue-600')

        if chart_type == '价格趋势':
            price_btn.classes(replace='bg-blue-700 text-white')
        elif chart_type == '平均收盘价对比':
            avg_btn.classes(replace='bg-blue-700 text-white')
        elif chart_type == '3D 散点图':
            scatter_btn.classes(replace='bg-blue-700 text-white')
        elif chart_type == '成交量饼图':
            pie_btn.classes(replace='bg-blue-700 text-white')

        chart_container.clear()
        with chart_container:
            if chart_type == '价格趋势':
                fig = create_price_trend_chart(current_filtered_df)
                ui.plotly(fig).classes('w-full h-[500px]')
            elif chart_type == '平均收盘价对比':
                fig = create_avg_close_price_bar_chart(current_filtered_df)
                ui.plotly(fig).classes('w-full h-[500px]')
            elif chart_type == '3D 散点图':
                fig = create_3d_scatter_chart(current_filtered_df)
                ui.plotly(fig).classes('w-full h-[500px]')
            elif chart_type == '成交量饼图':
                fig = create_volume_pie_chart(current_filtered_df)
                ui.plotly(fig).classes('w-full h-[500px]')

    def update_visualizations():
        """更新所有可视化组件"""
        nonlocal current_filtered_df

        # 清除现有内容
        kpi_container.clear()
        chart_container.clear()

        # 应用筛选条件
        selected_symbols = symbols.value
        # 从选择中移除特殊选项
        filtered_symbols = [s for s in selected_symbols if s not in special_options]

        # 如果选择了"所有代码"，则使用所有股票代码
        if '所有代码' in selected_symbols:
            filtered_symbols = all_symbols

        filtered_df = filter_data(
            df.copy(),
            start_date.value,
            end_date.value,
            filtered_symbols
        )

        # 保存当前筛选后的数据
        current_filtered_df = filtered_df

        if filtered_df.empty:
            # 显示提示信息而不是图表
            with chart_container:
                ui.label('请选择股票代码并点击"分析数据"按钮').classes('text-xl text-gray-500 mt-8')
            return

        # 创建 KPI 指标卡 - 改为对称布局
        with kpi_container:
            # 总记录数 - 蓝色卡片
            with ui.card().classes(
                    'flex-1 p-4 text-center bg-gradient-to-r from-blue-500 to-blue-700 text-white rounded-xl shadow-lg max-w-xs'):
                ui.label('总记录数').classes('text-lg font-bold')
                ui.label(f"{len(filtered_df):,}").classes('text-3xl font-bold mt-2')

            # 平均收盘价 - 绿色卡片
            avg_close = filtered_df['close'].mean()
            with ui.card().classes(
                    'flex-1 p-4 text-center bg-gradient-to-r from-green-500 to-green-700 text-white rounded-xl shadow-lg max-w-xs'):
                ui.label('平均收盘价').classes('text-lg font-bold')
                ui.label(f"${avg_close:.2f}").classes('text-3xl font-bold mt-2')

            # 总交易量 - 紫色卡片
            total_volume = filtered_df['volume'].sum()
            with ui.card().classes(
                    'flex-1 p-4 text-center bg-gradient-to-r from-purple-500 to-purple-700 text-white rounded-xl shadow-lg max-w-xs'):
                ui.label('总交易量').classes('text-lg font-bold')
                ui.label(f"{total_volume:,.0f}").classes('text-3xl font-bold mt-2')

        # 显示当前选中的图表
        switch_chart(current_chart_type)

        # 更新原始数据表格
        table_placeholder.clear()
        with table_placeholder:
            # 复制前 100 行用于显示
            df_display = filtered_df.head(100).copy()

            # 创建表格
            table = ui.table.from_pandas(df_display).classes('w-full max-h-80 overflow-auto')

    # 初始更新可视化
    update_visualizations()

# ------------------------------
# 程序入口
# ------------------------------
if __name__ in {"__main__", "__mp_main__"}:
    ui.run(title="股票数据分析平台", port=8081, dark=False)